'''
langchain agent

'''
from typing import Type

# from langchain import OpenAI, LLMChain
from langchain.agents import  initialize_agent, AgentType
import os

# 工具
# tools = load_tools(["tavily_search"], llm=None)  # 如果没有特别指定LLM，可以传入None

from langchain.tools import BaseTool
class MyCustomTool(BaseTool):
    name = "自定义工具"
    description = "自定义工具描述"

    def _run(self, query: str) -> str:
        # 在这里实现工具的主要逻辑
        return f"{query}"

    async def _arun(self, query: str) -> str:
        # 如果你想支持异步操作，实现这个方法
        raise NotImplementedError("This tool does not support async")

from pydantic import BaseModel, Field
class MyToolSchema(BaseModel):
    input_text: str = Field(..., description="The text to process.")
class MultiplicationTool(BaseTool):
    name = "糊糊糊工具"
    description = "自定义运算，运算名称为糊糊糊，输入两个整数，格式为 '10 12'，返回结果."
    args_schema: Type[BaseModel] = MyToolSchema
    def _run(self, query: str) -> str:
        try:
            num1, num2 = map(float, query.split())
            result = num1 * num2 + 100
            print(f" {num1} 乘 {num2} = {result}.")
            return f" {num1} 乘 {num2} = {result}."
        except ValueError:
            return f"输入格式错误 {query}"

# ty_tool
from langchain_community.tools import TavilySearchResults
os.environ["TAVILY_API_KEY"] = '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'
ty_tool = TavilySearchResults(
    max_results=5,
    search_depth="advanced",
    include_answer=True,
    include_raw_content=True,
    include_images=True,
    # include_domains=[...],
    # exclude_domains=[...],
    # name="...",            # overwrite default tool name
    description="一个为全面、准确和可信的结果而优化的搜索引擎。当你需要回答有关时事的问题时很有用。输入应该是一个搜索查询。"
    # args_schema=...,       # overwrite default args_schema: BaseModel
)
tools = [MyCustomTool(),MultiplicationTool(),ty_tool]


# 模型
from lc_infer import DeepSeek_R1_Distill_Qwen_LLM
# model_path = r'/home/ps/zhangxiancai/llm_deploy/bigfiles/models/DeepSeek-R1-Distill-Qwen-7B'
model_path = r'D:\code\other\LLMs\models\DeepSeek-R1-Distill-Qwen-1.5B'
llm = DeepSeek_R1_Distill_Qwen_LLM(mode_name_or_path=model_path)

# 代理
from langchain.agents import Agent
class CustomAgent(Agent):
    def _take_next_step(self, inputs):
        # 实现自定义逻辑
        print('实现自定义逻辑')
        pass
# agent = CustomAgent(llm=llm, tools=tools)

# agent_type = AgentType.SELF_ASK_WITH_SEARCH
agent_type = AgentType.ZERO_SHOT_REACT_DESCRIPTION # 是一种代理类型，它允许代理根据描述自动决定如何使用工具。
agent = initialize_agent(tools, llm, agent=agent_type, verbose=True)


# 运行
response = agent.run("今天上海有什么新闻")
print(response)